/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2021-2021.
 * Description: By using preload which directly load ko before starting the
 * init process. It's useful to make kernel smaller, Strip modules out of the
 * kernel, Support a log of boards to common mirror. Doesn't need switch_root.
 * Author: xiaowei
 * Create: 2021-08-31
 */

#include <linux/init.h>
#include <linux/vmalloc.h>
#include <linux/libfdt.h>
#include <linux/crc32.h>
#include <linux/memblock.h>
#include <linux/module.h>
#include <linux/decompress/generic.h>
#include <asm/sections.h>
#include <asm/setup.h>
#include <asm/errno.h>

#include "rtos_ko_initrd.h"

static unsigned int module_cnt;
static unsigned long preloadrd_start, preloadrd_end;
static phys_addr_t phys_preloadrd_start __initdata;
static unsigned long phys_preloadrd_size __initdata;

/* save boot_command_line for more lifetime */
char save_cmd_line[COMMAND_LINE_SIZE];
EXPORT_SYMBOL_NS(save_cmd_line, HW_RTOS_NS);

static __initdata char *message;
static void __init error(char *x)
{
	if (!message)
		message = x;
}

static int __init ko_preload_check_magic(struct ko_magic_hdr *header)
{
	unsigned int magic;

	magic = le32_to_cpu(header->magic);
	if (magic != KO_PRELOAD_MAGIC) {
		pr_err("check ko_magic_hdr fail. magic = 0x%x\n", magic);
		return -ENOEXEC;
	}

	return 0;
}

static void __init fill_ko_info(char *ipos, struct ko_info_hdr *info)
{
	struct ko_info_hdr *ko = NULL;

	ko = (struct ko_info_hdr *)ipos;
	(void)strlcpy(info->name, ko->name, KO_NAME_LEN);
	(void)strlcpy(info->args, ko->args, KO_ARGS_LEN);
	info->offset = le32_to_cpu(ko->offset);
	info->size = le32_to_cpu(ko->size);
}

static int __init rtos_unpack_preload_ko_packer(void *src, unsigned long src_size, unsigned long raw_size,
		void **dst, unsigned long *dst_size)
{
	int ret = 0;
	unsigned char *outbuf = NULL;
	decompress_fn decompress;
	const char *compress_name;
	message = NULL;

	if (src_size == 0) {
		pr_err("uncompress src_size ko_packer fail. please check hdr.vsize.\n");
		return -EINVAL;
	}

	decompress = decompress_method(src, src_size, &compress_name);
	pr_debug("Detected %s compressed data.\n", compress_name?compress_name:"NULL");
	if (!decompress) {
		pr_err("decompress failed.\n");
		return -EFAULT;
	}

	if (!raw_size) {
		pr_err("invalid packer size, must be nonzero.\n");
		return -EINVAL;
	}

	outbuf = vmalloc(raw_size);
	if (!outbuf)
		return -ENOMEM;

	ret = decompress(src, src_size, NULL, NULL, outbuf, NULL, error);
	if (ret) {
		pr_err("decompress failed ret=%d: %s.\n", ret, message);
		vfree(outbuf);
		return ret;
	}

	*dst = outbuf;
	*dst_size = raw_size;
	return 0;
}

static int __init rtos_load_preload_ko_modules(void *addr, unsigned int size)
{
	int ret = 0;
	int i;
	unsigned int cmfcrc;
	unsigned char *ipos = NULL;
	unsigned char *dpos = NULL;
	unsigned long unpack_size;
	void *unpack_addr = NULL;
	struct ko_magic_hdr *header = (struct ko_magic_hdr *)addr;
	struct ko_info_hdr info;

	ret = ko_preload_check_magic(header);
	if (ret)
		return ret;

	if (size < sizeof(struct ko_magic_hdr)) {
		pr_err("get header->vsize fail.\n");
		return -EFAULT;
	}

	cmfcrc = crc32(KO_CRC32_INIT, (unsigned char *)(addr + sizeof(struct ko_magic_hdr)), (size - sizeof(struct ko_magic_hdr)));
	if (header->crc != cmfcrc) {
		pr_err("CRC check fail: calculated %#08x, read %#08x.\n", cmfcrc, header->crc);
		return -EFAULT;
	}

	ret = rtos_unpack_preload_ko_packer((addr + sizeof(struct ko_magic_hdr)), (size - sizeof(struct ko_magic_hdr)), header->size, &unpack_addr, &unpack_size);
	if (ret)
		return ret;

	ipos = unpack_addr;
	dpos = unpack_addr + (ALIGN((module_cnt * KO_INFO_SIZE), PAGE_SIZE));

	for (i = 0; i < module_cnt; i++) {
		memset(&info, 0, sizeof(info));
		if ((ipos < (unsigned char *)unpack_addr) || ((dpos + info.offset) < (unsigned char *)unpack_addr) ||
			((ipos + KO_INFO_SIZE) > ((unsigned char *)unpack_addr + unpack_size)) ||
			((dpos + info.offset + info.size) > ((unsigned char *)unpack_addr + unpack_size))) {
			pr_err("load module size is more than vmalloc area.\n");
			vfree(unpack_addr);
			return -EFAULT;
		}

		fill_ko_info(ipos, &info);
		ret = ko_preload_load_module(info.name, dpos + info.offset, info.size, info.args);
		if (ret) {
			pr_err("load module fail.\n");
			vfree(unpack_addr);
			return ret;
		}
		pr_info("PRELOAD_OK:info.name info.offset info.size info.args = %s,%d,%d,%s.\n", info.name, info.offset, info.size, info.args);

		ipos += KO_INFO_SIZE;
	}

	vfree(unpack_addr);
	return 0;
}

static void __init arm_preloadrd_init(void)
{
	phys_addr_t start;
	unsigned long size;

	preloadrd_start = preloadrd_end = 0;
	if (!phys_preloadrd_size)
		return;

	start = round_down(phys_preloadrd_start, PAGE_SIZE);
	size = phys_preloadrd_size + (phys_preloadrd_start - start);
	size = round_up(size, PAGE_SIZE);

	if (!memblock_is_region_memory(start, size)) {
		pr_err("PRELOADRD: %pa+0x%08lx is not a memory region - disabling preloadrd.\n",
			&start, size);
		return;
	}

	if (memblock_is_region_reserved(start, size)) {
		pr_err("PRELOADRD: %pa+0x%08lx overlaps in-use memory region - disabling preloadrd.\n",
			&start, size);
		return;
	}

	memblock_reserve(start, size);

	strlcpy(save_cmd_line, boot_command_line, COMMAND_LINE_SIZE);

	preloadrd_start = __phys_to_virt(phys_preloadrd_start);
	preloadrd_end = preloadrd_start + phys_preloadrd_size;
}

static int __init early_preloadrd(char *p)
{
	phys_addr_t start;
	unsigned long size;
	char *endp;

	start = memparse(p, &endp);
	if (*endp == ',') {
		size = memparse(endp + 1, NULL);

		phys_preloadrd_start = start;
		phys_preloadrd_size = size;
	} else {
		pr_err("PRELOADRD: Unable to get the reserved memory size information.\n");
		return 0;
	}

	arm_preloadrd_init();

	return 0;
}
early_param("preloadrd", early_preloadrd);

static int __init rtos_ko_preload(void)
{
	int ret = 0;
	unsigned int vsize;
	unsigned int magic;
	struct ko_magic_hdr *header = NULL;
	void *preload_ko_addr = NULL;

	if (!preloadrd_start) {
		pr_err("PRELOADRD: 0x%lx res_mem can not resolve. please check the parameter of preloadrd in dts.\n",
			preloadrd_start);
		return -EFAULT;
	}
	header = (struct ko_magic_hdr *)preloadrd_start;
	pr_info("RTOS_KO_PRELOAD: begin to rtos_ko_preload!!!.\n");

	magic = le32_to_cpu(header->magic);
	if (magic != KO_PRELOAD_MAGIC) {
		pr_err("check header->magic fail, header->magic = 0x%x\n", magic);
		return -ENOEXEC;
	}

	module_cnt = le32_to_cpu(header->num);
	if ((!module_cnt) || (module_cnt > KO_NUM_MAX)) {
		pr_err("invalid, beyond[0,100].\n");
		return -EINVAL;
	}

	vsize = le32_to_cpu(header->vsize);
	if (!vsize)
		return -EINVAL;

	preload_ko_addr = vmalloc(vsize);
	if (!preload_ko_addr) {
		pr_err("vmalloc preload_ko space failed, header->vsize = %u\n", vsize);
		return -ENOMEM;
	}

	memcpy(preload_ko_addr, (void *)header, vsize);
	ret = rtos_load_preload_ko_modules(preload_ko_addr, vsize);
	if (ret) {
		pr_err("rtos_ko_preload load ko_packer fail.\n");
		vfree(preload_ko_addr);
		return ret;
	}
	vfree(preload_ko_addr);

	return 0;
}

module_init(rtos_ko_preload);
